﻿// Copyright (C) Jari Hautio <jari.hautio@iki.fi> 2008. Licensed under GPLv2. See LICENSE.txt in solution folder.
using System;
using System.Collections.Generic;
using System.Text;
using System.Drawing;

namespace ColorLib
{


    // Color cube based color optimizer
    public class CubePicker
    {
        CubeNode root = new CubeNode();
        List<Color> picked = new List<Color>();
        public enum CubeMode {Normal, LogCount,SqrtCount,Percent,MaxDistance, MaxError }
        CubeMode mode = CubeMode.Normal;
		public CubePicker(Dictionary<Color, int> cubeColors, CubeMode cubeMode)
		{
			mode = cubeMode;
			Init(cubeColors);
		}
		public CubePicker(ColorHistogram cubeColors, CubeMode cubeMode)
		{
			mode = cubeMode;
			Init(cubeColors);
		}
		/*
        Dictionary<Color, int> colors = new Dictionary<Color, int>();
        public Dictionary<Color, int> Colors
        {
            get { return colors; }
            set { colors = value; Init(); }
        }
		 */
		void Init(IEnumerable<KeyValuePair<Color, int>> colors)
        {
            root = new CubeNode();
            Color best = Color.Empty;
            int count = 0;
            int sum = 0;
            foreach (KeyValuePair<Color, int> t in colors)
            {
                if (count < t.Value)
                {
                    count = t.Value;
                    best = t.Key;
                }
                sum += t.Value;
            }
            picked.Add(best);
            foreach (KeyValuePair<Color, int> t in colors)
            {
                if (t.Key == best) continue;
                int n = t.Value;
                switch (mode)
                {
                    case CubeMode.LogCount:
                        n = Convert.ToInt32(Math.Log(t.Value)) + 1;
                        break;
                    case CubeMode.SqrtCount:
                        n = Convert.ToInt32(Math.Sqrt(1000000.0 * t.Value / count)) + 1;
                        break;
                    case CubeMode.Percent:
                        n = Convert.ToInt32(1000.0 * t.Value / count) + 1;
                        break;
                    case CubeMode.MaxDistance:
                        n = 1;
                        break;
                }
                root.Apply(new InsertVisitor(t.Key, best, n));
            }
        }
        // Make sure we have this many colors selected
        // and return selected colors
        public List<Color> Select(int n)
        {
            for (int i=picked.Count; i < n; i++)
            {
                SelectOne();
            }
            return picked;
        }
        Color SelectOne()
        {
            if (!root.Active) return Color.Empty;
            Color pickedColor = Color.Empty;
            PickVisitor v = null;
            switch (mode)
            {
                case CubeMode.MaxDistance:
                    v = new PickDistanceVisitor();
                    break;
                case CubeMode.MaxError:
                    v = new PickMaxErrorVisitor();
                    break;
                default:
                     v = new PickVisitor();
                    break;
            }
            root.Apply(v);
            if (v.Picked != null)
            {
                pickedColor = v.Picked.color;
                picked.Add(pickedColor);
            }

            return pickedColor;
        }
        public Dictionary<Color,Color> GetColorMap()
        {
            ColorMapVisitor v = new ColorMapVisitor();
            root.Apply(v);
            return v.ColorMap;
        }
        public CubeNode Cube { get { return root; } }

        public static double MaxErrorValue = 4 * 255 * 255;
        public class CubeVisitor
        {
            int maxDepth = 7;
            int minDepth = 2;
            public void Visit(CubeNode c) 
            {
                level++;
                nodePath.AddLast(c);
                OnVisit(c);
                nodePath.RemoveLast();
                level--;
            }
            protected void VisitAllNodes(CubeNode c)
            {
                foreach (ICubeNode n in c.Slots)
                {
                    if (n != null)
                        n.Apply(this);
                }
            }
            protected void VisitTargetNode(CubeNode c)
            {
                int s = c.SelectNode(Target);
                ICubeNode nextNode = c.Slots[s];
                // if not yet created, then create new node
                if (nextNode == null)
                {
                    if (AllowCreate)
                    {
                        if (c.Level < minDepth)
                        {
                            c.CreateChild(s);
                        }
                        else
                        {
                            c.CreateLeaf(s);                            
                        }
                    }
                    else // not found
                    {
                        return;
                    }
                }
                else if (c.Slots[s] is ColorNode)
                {
                    ColorNode colorNode = nextNode as ColorNode;
                    if (colorNode.Center != Target)
                    {
                        if (AllowCreate && c.Level < maxDepth) c.AddLevel(s);
                        else return; // not found
                    }
                    // else ok to go 

                }
                // else it's a cube node on next level
                c.Slots[s].Apply(this);
            }
            protected virtual void OnVisit(CubeNode c)
            {
                // visit all if target color is empty
                if (Target == Color.Empty)
                {
                    VisitAllNodes(c);
                }
                else // head for target
                {
                    VisitTargetNode(c);
                }
            }
            public void Visit(ColorNode c)
            {
                level++;
                nodePath.AddLast(c);
                OnVisit(c);
                nodePath.RemoveLast();
                level--;
            }
            protected virtual void OnVisit(ColorNode c)
            {
            }
            Color target = Color.Empty;
            public  Color Target { get { return target; } set { target = value; } }
            bool allowCreate = false;
            public bool AllowCreate { get { return allowCreate; } set { allowCreate = value; } }
            private LinkedList<ICubeNode> nodePath = new LinkedList<ICubeNode>();
            protected CubeNode ParentNode { get { return nodePath.Count > 1 ?  (nodePath.Last.Previous.Value as CubeNode): null; } }
            public LinkedList<ICubeNode> NodePath { get { return nodePath; } }
            public int Level { get { return level;}}
            public double CalcDistance(Color c1, Color c2)
            {
                return 0.5*Math.Sqrt(ColorLib.ColorMath.DistanceYCbCrFast(c1,c2))+
                        0.5*Math.Sqrt((c1.R - c2.R) * (c1.R - c2.R) + (c1.G - c2.G) * (c1.G - c2.G) + (c1.B - c2.B) * (c1.B - c2.B));
            }
            public double ErrorFunc(Color target, double distance, int pixelCount)
            {
                return distance * distance * pixelCount;
            }
            int level=0;

        }
        public class ColorMapVisitor : CubeVisitor
        {
            Dictionary<Color, Color> map = new Dictionary<Color,Color>();
            public Dictionary<Color, Color> ColorMap { get { return map; } }
            protected override void OnVisit(ColorNode c)
            {
                map.Add(c.data.color, c.data.nearest);
            }

        }
        public class InsertVisitor : CubeVisitor
        {            
            double oldError = 0;
            double newError = 0;
            ColorData data = new ColorData();
            ColorNode insertionNode = null;
            public ColorNode InsertionNode { get { return insertionNode; } }
            public InsertVisitor(Color target, double distance, Color nearestColor, int pixelCount)
            {
                Target = target;
                AllowCreate = true;
                data.color = target;
                data.count = pixelCount;
                data.nearest = nearestColor;
                data.distance = distance;
                data.error = ErrorFunc(target, distance ,pixelCount);
                
            }
            public InsertVisitor(Color target, Color nearestColor, int pixelCount)
            {
                Target = target;
                AllowCreate = true;
                data.color = target;
                data.count = pixelCount;
                data.nearest = nearestColor;
                data.distance = CalcDistance(target, nearestColor);
                data.error = ErrorFunc(target, data.distance ,pixelCount);

            }
            protected override void OnVisit(ColorNode cn)
            {
                oldError = cn.data.error;
                cn.data = data;
                newError  = data.error;
                insertionNode = cn;
            }
            protected override void OnVisit(CubeNode c)
            {
                base.OnVisit(c);
                // update error values on way back
                // set max distance to max distance of children
                // TODO optimize this!
                c.sumError = 0;
                c.maxDistance = 0;
                c.maxError = 0;
                foreach (ICubeNode n in c.Slots)
                {
                    if (n == null) continue;
                    if (c.maxDistance < n.Distance) c.maxDistance = n.Distance;
                    if (c.maxError < n.MaxError) c.maxError = n.MaxError;
                    c.sumError += n.SumError;
                }
                c.UpdateBound(insertionNode);
                c.Count += insertionNode.Count;
            }
        }

        // Update error and distance values for each node against target color
        public class ErrorVisitor : CubeVisitor
        {
            protected override void OnVisit(ColorNode n)
            {
                if (!n.Active) return;
                Color c = n.data.color;
                double d = CalcDistance(c, Target);
                // if distance is larger than previous distance, then target
                // color wont affect to this node
                if (d >= n.data.distance)
                    return;
                // Recalculate error 
                n.data.distance = d;
                n.data.error = ErrorFunc(c,d,n.data.count);
                n.data.nearest = Target;
            }
            protected override void OnVisit(CubeNode c)
            {
                /*
                 * old code for version with spherical bounds
                double d = CalcDistance(Target, c.Center) - c.Width *0.867;
                if (d > c.Distance)
                   return;
                 * */
                if (c.OutsideBounds(Target))
                    return;
                c.Distance = 0;
                c.sumError = 0;
                c.maxError = 0;
                c.NoBound = true;
                foreach (ICubeNode n in c.Slots)
                {
                    if (n == null) continue;
                    n.Apply(this);
                    if (c.Distance < n.Distance) c.Distance = n.Distance;
                    if (c.MaxError < n.MaxError) c.maxError = n.MaxError;
                    c.sumError += n.SumError;
                    c.UpdateBound(n);
                }
            }
        }
        public class DumpVisitor : CubeVisitor
        {
            StringBuilder sb = new StringBuilder();
            string path;
            public string Dump { get { return sb.ToString(); } }
            public DumpVisitor()
            {
                sb.Append("Type\tLvl\tPath\tCount\tDist\tMaxError\tSumError\n");
            }
            
            protected override void OnVisit(CubeNode c)
            {
                string s = String.Format("Node\t  {0}\t{1}\t{2}\t{3}\t{4}\t{5}\t{6}\t{7}\t{8}\n",
                            Level, path, c.Count,c.Distance, c.MaxError,c.SumError,
                            c.Center.R, c.Center.G, c.Center.B);
                sb.Append(s);
                int t=0;
                string oldpath = path;
                foreach (ICubeNode n in c.Slots)
                {
                    path = oldpath + t.ToString();
                    if (n != null)
                        n.Apply(this);
                    t++;
                }
                path = oldpath;
            }
            protected override void OnVisit(ColorNode c)
            {
                string s = String.Format("Leaf\t  {0}\t{1}\t{2}\t{3}\t{4}\t{5}\t{6}\t{7}\t{8}\n",
                            Level, path, c.Count,c.Distance, c.MaxError,c.SumError,
                            c.Center.R, c.Center.G, c.Center.B);
                sb.Append(s);
            }
        }
        /// <summary>
        /// Picks a color with largest error value so that 
        /// the sum of error in each level of path to target color
        /// is largest.
        /// </summary>
        public class PickVisitor : CubeVisitor
        {
            ColorData picked = null;
			// Setter provided for testing
			public ColorData Picked { get { return picked; } set { picked = value; } }
            protected virtual ICubeNode SelectNode(CubeNode c)
            {
                ICubeNode best = null;
                int tmp = c.Count > 0 ? c.Count : 1;// (0x80 >> c.Level);
                double bestDist=0;
                foreach (ICubeNode node in c.Slots)
                {
                    if (node == null) continue;
                    double d = Math.Pow(node.Distance, 3) + node.MaxError + node.SumError / tmp;
                    if (best == null)
                    {
                        best = node;
                        bestDist = d;
                    }
                    else if (bestDist < d)
                    {
                        best = node;
                        bestDist = d;
                    }
                }
                return best;
            }
            // select always node with largest error sum
            protected override void OnVisit(CubeNode c)
            {
                ICubeNode best= SelectNode(c);
                if (best != null)
                {
                    // Enter next level
                    best.Apply(this);
                    if (picked == null)
                        return;
                    // Update node state
                    if (c.Active)
                    {
                        UpdateNode(c, best);
                        c.Count -= picked.count;
                    }
                    else
                    {
                        if (ParentNode != null)
                            ParentNode.RemoveChild(c);
                        else
                        {
                            c.sumError = 0;
                            c.maxError = 0;
                            c.maxDistance = 0;
                            c.NoBound = true;
                        }
                    }
                }
            }

            protected void UpdateNode(CubeNode c, ICubeNode skip)
            {
                if (picked == null) return;
                // update distance and error values
                ErrorVisitor v = new ErrorVisitor();
                v.Target = picked.color;
                c.Distance = 0;
                c.sumError = 0;
                c.maxError = 0;
                c.NoBound = true;
                foreach (ICubeNode n in c.Slots)
                {
                    if (n == null ) continue;
                    if (n != skip) n.Apply(v);
                    if (c.Distance < n.Distance) c.Distance = n.Distance;
                    if (c.MaxError < n.MaxError) c.maxError = n.MaxError;
                    c.sumError += n.SumError;
                    c.UpdateBound(n);
                }
            }
            protected override void OnVisit(ColorNode c)
            {
                // found color candidate
                picked = c.data;
                // remove from tree
                ParentNode.RemoveChild(c);
            }

        }

        /// <summary>
        /// Picks a color with largest error value so that 
        /// the sum of error in each level of path to target color
        /// is largest.
        /// </summary>
        public class PickDistanceVisitor : PickVisitor
        {
            // select always node with largest distance to a picked color
            protected override ICubeNode SelectNode(CubeNode c)
            {
                ICubeNode best = null;
                foreach (ICubeNode n in c.Slots)
                {
                    if (n == null) continue;
                    if (best == null)
                        best = n;
                    else if (best.Distance < n.Distance)
                        best = n;
                }
                return best;
            }

        }
        /// <summary>
        /// Picks a color with largest error value so that 
        /// the sum of error in each level of path to target color
        /// is largest.
        /// </summary>
        public class PickMaxErrorVisitor : PickVisitor
        {
            // select always node with largest error sum
            protected override ICubeNode SelectNode(CubeNode c)
            {
                ICubeNode best = null;
                foreach (ICubeNode n in c.Slots)
                {
                    if (n == null) continue;
                    if (best == null)
                        best = n;
                    else if (best.MaxError < n.MaxError)
                        best = n;
                }
                return best;
            }

        }
        
        public interface ICubeNode
        {
            void Apply(CubeVisitor v);
            double MaxError { get; }
            double SumError { get; }
            double Distance { get; }
            bool Active { get; }
            int Width { get; }
            Color Center { get; }
            byte MinR { get; }
            byte MinG { get; }
            byte MinB { get; }
            byte MaxR { get; }
            byte MaxG { get; }
            byte MaxB { get; }
            bool WithinBounds(Color c);
            int Count { get; set; }
        }
        public class ColorData
        {
            public ColorData()
            {
            }
            public ColorData(Color c, int n, double e, double d, Color nc)
            {
                color = c;
                count = n;
                error = e;
                distance = d;
                nearest = nc;
            }
            public Color color = Color.Empty;
            public int count = 0;
            public double error = 0;
            public double distance = 0;
            // nearest picked color to which error is calculated
            public Color nearest = Color.Empty;

        }
        public class ColorNode : ICubeNode
        {
            public ColorNode(ColorData cd) { data = cd; }
            public ColorNode() { data = new ColorData(); }
            public bool Active { get { return data.color != Color.Empty; } }
            public ColorData data = new ColorData();
            public double MaxError { get { return data.error; } }
            public double SumError { get { return data.error; } }
            public double Distance { get { return data.distance; } }
            public int Count { get { return data.count; } set { data.count = value; } }
            public void Apply(CubeVisitor v)
            {
                v.Visit(this);
            }
            public int Width { get { return 1; } }
            public Color Center { get { return data.color; } }
            public byte MinR { get { return ToByte(data.color.R - data.distance,false); } }
            public byte MinG { get { return ToByte(data.color.G - data.distance,false); } }
            public byte MinB { get { return ToByte(data.color.B - data.distance,false); } }
            public byte MaxR { get { return ToByte(data.color.R + data.distance,true); ; } }
            public byte MaxG { get { return ToByte(data.color.G + data.distance,true); } }
            public byte MaxB { get { return ToByte(data.color.B + data.distance,true); } }
            byte ToByte(double d, bool ceil)
            {
                int n = Convert.ToInt32(ceil ? Math.Ceiling(d) : Math.Floor(d));
                if (d < 0) return 0;
                if (d > 255) return 255;
                return (byte)n;
            }
            public bool WithinBounds(Color c)
            {
                if (!Active) return false;
                return MinR <= c.R &&
                        MaxR >= c.R &&
                        MinG <= c.G &&
                        MaxG >= c.G &&
                        MinB <= c.B &&
                        MaxB >= c.B;
            }
        }
        public class CubeNode : ICubeNode
        {
            public static int nodeCount=0;
//            int maxDepth = 7;
            int level = 0;
            int count=0;
            byte used = 0;
            Color centerColor = Color.Empty;
            protected ICubeNode[] slots = new ICubeNode[8];
            public ICubeNode[] Slots { get { return slots; } }
            public byte minR,minG,minB,maxR,maxG,maxB;
            public double sumError = 0;
            public double maxError = 0;
            public double SumError { get { return sumError; } }
            public double MaxError { get { return maxError; } }
            public int Level { get { return level; } }
            public double Distance { get { return maxDistance; } set { maxDistance = value; } }
            public int Count { get { return count; } set { count = value; } }
            /// <summary>
            /// Max Distance from node center to farthest color that
            /// is currently used as nearest color
            /// </summary>
            public double maxDistance = 0;
            public bool Active { get { return used > 0; } }
            public byte MinR { get { return minR; } }
            public byte MinG { get { return minG; } }
            public byte MinB { get { return minB; } }
            public byte MaxR { get { return maxR; } }
            public byte MaxG { get { return maxG; } }
            public byte MaxB { get { return maxB; } }
            bool noBound = true;
            public bool NoBound { get { return noBound; } set { noBound = value; } }

            public CubeNode()
            {
                level = 0;
                centerColor = Color.FromArgb(128, 128, 128);
                nodeCount++;
            }
            public CubeNode(int nodeLevel, Color center)
            {
                level = nodeLevel;
                centerColor = center;
                nodeCount++;
            }
            public void UpdateBound(Color c)
            {
                if (noBound || MinR > c.R) minR = c.R;
                if (noBound || MaxR < c.R) maxR = c.R;
                if (noBound || MinG > c.G) minG = c.G;
                if (noBound || MaxG < c.G) maxG = c.G;
                if (noBound || MinB > c.B) minB = c.B;
                if (noBound || MaxB < c.B) maxB = c.B;
                noBound = false;
            }
            public void UpdateBound(ICubeNode c)
            {
                if (noBound || MinR > c.MinR) minR = c.MinR;
                if (noBound || MaxR < c.MaxR) maxR = c.MaxR;
                if (noBound || MinG > c.MinG) minG = c.MinG;
                if (noBound || MaxG < c.MaxG) maxG = c.MaxG;
                if (noBound || MinB > c.MinB) minB = c.MinB;
                if (noBound || MaxB < c.MaxB) maxB = c.MaxB;
                noBound = false;
            }
            public bool WithinBounds(Color c)
            {
                if (noBound) return false;
                return  minR <= c.R &&
                        maxR >= c.R &&
                        minG <= c.G &&
                        maxG >= c.G &&
                        minB <= c.B &&
                        maxB >= c.B;
            }
            public bool OutsideBounds(Color c)
            {
                if (noBound) return true;
                return  (minR > c.R || maxR < c.R) ||
                        (minG > c.G || maxG < c.G) ||
                        (minB > c.B || maxB < c.B);
            }
            public CubeNode CreateChild(int slot)
            {
                if (slots[slot] != null) return slots[slot] as CubeNode;

                int w = Width / 4;
                Color center = Color.FromArgb(
                    centerColor.R + (((slot & 1) != 0) ? w : -w),
                    centerColor.G + (((slot & 2) != 0) ? w : -w),
                    centerColor.B + (((slot & 4) != 0) ? w : -w));
                CubeNode newNode = new CubeNode(Level + 1, center);
                slots[slot] = newNode;
                used++;
                return newNode;
            }
            public ICubeNode CreateLeaf(int slot)
            {
                if (slots[slot] != null) return slots[slot] ;
                slots[slot] = new ColorNode();
                used++;
                return slots[slot];
            }
            public void Apply(CubeVisitor v)
            {
                v.Visit(this);

            }
            public int SelectNode(Color c)
            {
                return  1 * (((c.R << level) & 0x80) >> 7) +
                        2 * (((c.G << level) & 0x80) >> 7) +
                        4 * (((c.B << level) & 0x80) >> 7);
            }
            public void RemoveChild(ICubeNode n)
            {
                if (n == null) return;
                for (int i=0;i<slots.Length;i++)
                {
                    if (slots[i] == n)
                    {
                        slots[i] = null;
                        used--;
                    }
                }
            }

            public CubeNode AddLevel(int slot)
            {                
                if (slots[slot] == null)
                    return CreateChild(slot);
                ColorNode cn = slots[slot] as ColorNode;
                if (cn == null)
                {
                    return slots[slot] as CubeNode; // already a cube node
                }
                // replace color node with cube node and reinsert color node
                slots[slot] = null;
                used--;
                CubeNode newNode = CreateChild(slot);
                int newSlot = newNode.SelectNode(cn.data.color);
                newNode.Slots[newSlot] = cn;
                newNode.Count = cn.Count;
                newNode.Distance = cn.Distance;
                newNode.maxError = cn.MaxError;
                newNode.sumError = cn.SumError;
                newNode.UpdateBound(cn);
                newNode.used = 1;
                return newNode;
            }
            // Level 0 --> 256, Level 7 --> 1
            public int Width { get { return 256 >> Level; } }
            public Color Center { get { return centerColor; } }
        }
    }
}
